# Copyright (c) 2024 刻BITTER
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.


from typing import Optional, Dict, List, Any, Union
import json
from pathlib import Path
import os
from itertools import chain

import rock_hammer as rh


def restore_rpgmvp(f: Path, key_str: str) -> Optional[Path]:
    '''
    png 文件的前16 字节被加密，然后加上了16 字节rpgmvp 的文件头，也就是文件开头的32 个字节部分。
    还原的原理就是干脆把这32 字节剪切掉，然后把正常的png 文件前面16 字节部分拼接上去
    '''
    assert f.exists()
    if f.suffix not in ('.rpgmvp', '.png_'):
        print(rh.red(f'!-> Invalid RPGMVP file: {str(f)}'))
        return None
    png_header = bytes(b'\x89\x50\x4E\x47\x0D\x0A\x1A\x0A\x00\x00\x00\x0D\x49\x48\x44\x52')
    header_len = len(png_header)
    dat = f.read_bytes()
    dat = dat[header_len * 2 :]
    restore_dat = png_header + dat
    png_name = f.stem + '.png'
    new_file = f.parent / png_name
    new_file.write_bytes(restore_dat)
    return new_file


def get_key_str(json_dict: dict[str, Any]) -> str:
    assert 'encryptionKey' in json_dict
    key_str = json_dict['encryptionKey']
    assert isinstance(key_str, str) and len(key_str) % 2 == 0
    return key_str


def restore_rpgmvo(f: Path, key_str: str) -> Optional[Path]:
    assert f.exists()
    if f.suffix not in ('.rpgmvo', '.ogg_'):
        print(rh.red(f'!-> Invalid RPGMVO file: {str(f)}'))
        return None

    '''
    解密原理是，先文件头部header_len 长度的部分截掉，这是加密是添加的文件头，
    之后是被加密的原始文件数据，加密的部分只有头部header_len 个数的字节，
    并且与加密key  的每个字节对应，只需用加密key 对其异或复原即可
    '''
    key_array = bytes.fromhex(key_str)
    header_len = len(key_array)
    assert header_len > 0
    dat = f.read_bytes()
    dat = dat[header_len:]
    revert = []
    for i in range(header_len):
        d = dat[i]
        revert.append(d ^ key_array[i])
    restore_dat = bytes(revert) + dat[header_len:]
    ogg_name = f.stem + '.ogg'
    new_file = f.parent / ogg_name
    new_file.write_bytes(restore_dat)
    return new_file


def get_game_sub_path(name: str, game_dir: Path) -> Optional[Path]:
    base_path = game_dir / 'www'
    if not base_path.exists():
        base_path = game_dir
    
    if name == 'img' or name == 'audio':
        path = base_path
    else:
        path = base_path / name
    
    if not path.exists():
        return None
    else:
        return path


def restore_file(game_folder: Path, json_dict: dict[str, Any], *, file_type:str, delete_encrypted_file=False) -> None:
    file_type = file_type.lower()
    if file_type == 'image':
        json_flag_name = 'hasEncryptedImages'
        dir_name = 'img'
        ext = ['rpgmvp', 'png_']
        revert_function = restore_rpgmvp
    elif file_type == 'audio':
        json_flag_name = 'hasEncryptedAudio'
        dir_name = 'audio'
        ext = ['rpgmvo', 'ogg_']
        revert_function = restore_rpgmvo
    else:
        raise ValueError(f'Unsupported File Type: {file_type}')

    if json_flag_name in json_dict:
        dir_path = get_game_sub_path(dir_name, game_folder)
        assert dir_path is not None and dir_path.is_dir()
        print(f'-> Start Decryption: {file_type}')

        count = 0
        key_str = get_key_str(json_dict)
        for f in chain(*[dir_path.rglob(f'*.{e}') for e in ext]):
            nf = revert_function(f, key_str)
            if nf is not None and delete_encrypted_file:
                os.remove(f)
            count += 1

        print(f'-> Decrypted File Count: [ {count} ]')
        json_dict[json_flag_name] = False


def get_system_json(game_folder: Path) -> dict[str, Any]:
    json_path = get_game_sub_path('data/system.json', game_folder)
    assert json_path is not None and json_path.is_file()
    json_text = json_path.read_text(encoding='utf-8')
    r = json.loads(json_text)
    assert isinstance(r, dict)
    return r


def write_back_system_json(game_folder: Path, json_dict: dict):
    json_path = get_game_sub_path('data/system.json', game_folder)
    assert json_path is not None and json_path.is_file()
    text = json.dumps(json_dict)
    json_path.write_text(text, encoding='utf-8')


def restore_img_and_audio(game_dir: Union[Path, str], *, delete_encrypted_file=False, silent=False):
    game_dir = Path(game_dir)

    if delete_encrypted_file:
        rh.sys_print(f'?> Original File Will Be Deleted')
        if not silent:
            rh.sys_print(f'?> Confirm?')
            input()
    j = get_system_json(game_dir)
    restore_file(game_dir, j, file_type='image', delete_encrypted_file=delete_encrypted_file)
    restore_file(game_dir, j, file_type='audio', delete_encrypted_file=delete_encrypted_file)
    write_back_system_json(game_dir, j)


if __name__ == "__main__":
    d = Path(r"D:\The Slave Race Elf and the Female Body Armor Ogre")
    restore_img_and_audio(d, delete_encrypted_file=True, silent=True)
